import random
from IPython.core.display import SVG
import pyomo.environ as pyo
from pysat.solvers import Solver
from pysat.formula import CNF 
import py_svg_combinatorics as psc
from ipywidgets import widgets, HBox
from collections import Counter
from pprint import pprint
from random import randint
import numpy as np
from IPython.display import IFrame
import IPython
from copy import copy
import os
from pathlib import Path
from pyomo_helpers import pretty_solution
nbname = ''
try:
    nbname = __vsc_ipynb_file__
except:
    if 'COCALC_JUPYTER_FILENAME' in os.environ:
        nbname = os.environ['COCALC_JUPYTER_FILENAME']    
title_ = Path(nbname).stem.replace('-', '_').title()    
IFrame(f'https://discopal.ispras.ru/index.php?title=Hardprob/{title_}&useskin=cleanmonobook', width=1280, height=300)

#phys_regs_amount = 3 — слишком скучный вариант
phys_regs_amount = 4
phys_regs_costs = np.arange(1, phys_regs_amount + 1)

basic_block_size = 16
instruction_uses_max = int(phys_regs_amount / 2)  # при phys_regs_amount = 3 — тут будет 1, слишком скучный вариант
#instruction_uses_max = int(phys_regs_amount * 3 / 4)  # при phys_regs_amount = 3 — тут будет 1, слишком скучный вариант

virtual_alives = [0, 1]
virtual_uses = []

for _ in range(basic_block_size):
    # Каждая инструкция использует до instruction_uses_max виртуальных регистров
    virtual_uses.append([f"{i:02}" for i in np.random.choice(virtual_alives, instruction_uses_max)])
    
    # Инструкция с шансом 0.5 может определить новый виртуальный регистр
    if np.random.choice([True, False]):
        virtual_alives.append(virtual_alives[-1] + 1)

virtual_uses

svg = psc.subsets2svg(virtual_uses)
SVG(data=svg)

def get_model(instruction_uses, phys_regs_costs, virtual_regs_amount):
    model = pyo.ConcreteModel()
    
    # Множество физических регистров
    model.phys_regs = pyo.Set(initialize=range(len(phys_regs_costs)))
    
    # Сопоставление каждому физическому регистру его стоимости  -- не совсем понятно, зачем нам параметры в конкретной модели
    model.phys_costs = pyo.Param(model.phys_regs, initialize=phys_regs_costs)
    
    # Множество виртуальных регистров
    model.virt_regs = pyo.Set(initialize=range(virtual_regs_amount))
    
    # Хотим составить отображение из виртуальных регистров на физические
    model.alloca = pyo.Var(model.virt_regs, model.phys_regs, domain=pyo.Binary)
    
    # Составляем множества плохих пар виртуальных регистров (встречаемых в используемых для одинаковых инструкций)
    model.bad_pairs = pyo.Set(  # !!! плохой пример — получается пустое множество. ij — символы, очень мутно
        initialize=sorted(set([(int(i), int(j)) for uses in instruction_uses for i in uses for j in uses if i < j])))
    
    # ЦЕЛЬ: минимизировать стоимость использования физических регистров после аллокации
    model.obj = pyo.Objective(expr=sum(model.phys_costs[p] * model.alloca[v, p] 
                                for v in model.virt_regs 
                                for p in model.phys_regs))
    
    # ОГРАНИЧЕНИЕ: каждому виртуальному регистру сопоставляем ровно 1 физический
    @model.Constraint(model.virt_regs)
    def виртуальному_один_физический(m, v):
        return sum(model.alloca[v, phys] for phys in m.phys_regs) == 1
    
    # ОГРАНИЧЕНИЕ: каждый физический регистр должен использоваться каждой инструкцией не более одного раза
    @model.Constraint(model.bad_pairs, model.phys_regs)
    def использование_одного_физического(m, first, second, reg):
        return (m.alloca[first, reg] + m.alloca[second, reg]) <= 1
    
    # @model.Constraint(range(8))
    # def wtf(m, в):
    #     if в==0:
    #         return model.alloca[в, 1] == 1
    #     elif в==1:
    #         return model.alloca[в, 0] == 1
    #     elif в==2:
    #         return model.alloca[в, 2] == 1
    #     elif в==3:
    #         return model.alloca[в, 0] == 1
    #     elif в==4:
    #         return model.alloca[в, 3] == 1
    #     elif в==5:
    #         return model.alloca[в, 1] == 1
    #     elif в==6:
    #         return model.alloca[в, 0] == 1
    #     elif в==7:
    #         return model.alloca[в, 0] == 1
    #     return  pyo.Constraint.Skip
    
    return model
    

m = get_model(virtual_uses, phys_regs_costs, len(virtual_alives))
m.pprint()

solver = pyo.SolverFactory('cbc')
solver.solve(m).write()

#m.alloca.extract_values()
pretty_solution(m, skip_zeros=True)

for v in m.component_data_objects(pyo.Var):
    if v.value and v.value > 0:
        print(str(v), v.value)

def модель(инструкции, стоимость_регистров):
    m = pyo.ConcreteModel()
    m.В  = sorted(set([в for и in инструкции for в in и]))
    m.Ф  = range(len(стоимость_регистров))
    m.I  = range(len(инструкции))
    m.вирт2физ = pyo.Var(m.В, m.Ф, domain=pyo.Binary)
    
    @m.Objective()
    def стоимость_аренды(m, в):
        return sum(стоимость_регистров[ф] * m.вирт2физ[в, ф] for в in m.В for ф in m.Ф)
    
    @m.Constraint(m.В)
    def виртуальному_один_физический(m, в):
        return sum(m.вирт2физ[в, ...]) == 1

    @m.Constraint(m.Ф, m.I)
    def в_каждой_инструкции_мапировать_не_больше_чем_в_один_физический(m, ф, i):
        вирт_регистры = sorted(set(инструкции[i]))
        выражение = 0 
        for имя_регистра in вирт_регистры:
            выражение += m.вирт2физ[имя_регистра, ф]
        return выражение <= 1
   
    return m

m2 = модель(virtual_uses, phys_regs_costs)
#m2.pprint()
solver = pyo.SolverFactory('cbc')
solver.solve(m2).write()

virtual_uses

pretty_solution(m2, skip_zeros=True)

phys_uses = []
for use in virtual_uses:
    phys_uses.append([f"{p:02}" for v in use for p in range(phys_regs_amount) if pyo.value(m.alloca[int(v), p])])
svg = psc.subsets2svg(phys_uses)
SVG(data=svg)

sat_vars_amount = 5
sat_clauses_amount = 20

cnf3 = CNF(from_clauses=psc.rand3cnf(sat_clauses_amount, sat_vars_amount))
cnf3.clauses

solver = Solver(bootstrap_with=cnf3)
res = solver.solve()
res

NEUTRAL = 0
TRUE = 1
FALSE = 2

vertexes_amount = 3
edges = [(NEUTRAL, TRUE), (TRUE, FALSE), (FALSE, NEUTRAL)]

def create_var():
    global vertexes_amount
    
    new_var = vertexes_amount
    vertexes_amount += 1

    return new_var

def ident_var():
    global edges

    new_var = create_var()
    edges.append((new_var, NEUTRAL))

    return new_var

def negate_var(to_negate):
    global edges
    
    negative_var = ident_var()
    edges.append((negative_var, to_negate))

    return negative_var

def or_var(lhs, rhs):
    global edges
    global vertexes_amount

    res_var = ident_var()

    tmps = list(range(vertexes_amount, vertexes_amount + 4))
    vertexes_amount += 4

    edges.append((TRUE, tmps[0]))

    edges.append((tmps[0], tmps[1]))
    edges.append((tmps[0], res_var))

    edges.append((tmps[1], lhs))
    edges.append((tmps[1], rhs))

    edges.append((lhs, tmps[2]))
    edges.append((rhs, tmps[3]))

    edges.append((tmps[2], tmps[3]))
    edges.append((tmps[2], res_var))
    edges.append((tmps[3], res_var))

    return res_var

import functools

var_converter = {}
for sat_var in range(1, sat_vars_amount + 1):
    var_converter[sat_var] = ident_var()

negate_converter = {}
for sat_var in range(1, sat_vars_amount + 1):
    negate_converter[sat_var] = negate_var(var_converter[sat_var])

converted_clauses = []
for clause in cnf3.clauses:
    converted_clauses.append(functools.reduce(or_var, map(lambda x: var_converter[x] if x > 0 else negate_converter[-x], clause)))

for clause in converted_clauses:
    edges.append((clause, NEUTRAL))
    edges.append((clause, FALSE))

m = get_model(edges, phys_regs_costs, vertexes_amount)
m.pprint()

cbc_solver = pyo.SolverFactory('cbc')
cbc_solver.solve(m).write()

false_reg = [p for p in range(phys_regs_amount) if pyo.value(m.alloca[FALSE, p])]
true_reg = [p for p in range(phys_regs_amount) if pyo.value(m.alloca[TRUE, p])]

false_reg = false_reg[0]
true_reg = true_reg[0]

answer = []
for sat_var in range(1, sat_vars_amount + 1):
    is_false = pyo.value(m.alloca[var_converter[sat_var], false_reg])
    is_true = pyo.value(m.alloca[var_converter[sat_var], true_reg])
    
    assert(is_false or is_true)
    if is_true:
        answer.append(sat_var)
    else:
        answer.append(-sat_var)

print(answer)
print(list(solver.enum_models()))